package com.webshh.websocket;

import com.jcraft.jsch.Channel;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.UserInfo;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.stereotype.Service;
import org.springframework.web.socket.server.standard.SpringConfigurator;

import javax.annotation.PostConstruct;
import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.ByteBuffer;
import java.util.concurrent.CopyOnWriteArraySet;

/**
 * @author frank
 */
@ServerEndpoint(value = "/webshell/{host}/{ip}")
@Component
public class ShellSocket {
    private Logger logger = LoggerFactory.getLogger(getClass());
    private final String CHAESET_NAME = "utf-8";

    public final String THREAD_NAME = "JSCH-READER:";

    private final String STRICT_HOST_KEY_CHECKING = "StrictHostKeyChecking";
    private final String NO = "no";
    private final String JSCH_TYPE = "shell";
    private final int SESSION_TIMEOUT = 30000;
    private final int CHANNEL_TIMEOUT = 3000;

    private static final String USER = "root";
    private static final String PASSWORD = "BsdldSsh123456";
    private static String HOST;
    private static final int DEFAULT_SSH_PORT = 22;

    private static CopyOnWriteArraySet<String> webSocketSet = new CopyOnWriteArraySet<String>();

    private Session session;
    IoReader reader = null;
    private Channel channel;
    private OutputStream out;
    private InputStream in;
    private String userIp;

    /**
     * 连接建立成功调用的方法
     */
    @OnOpen
    public void onOpen(@PathParam("host") String host,@PathParam("ip") String ip, Session webShell) throws JSchException {
        if (webSocketSet.contains(ip)) {
            return;
        }
        System.out.println("host = " + host + "-------------------------------------- ip = " + ip);
        HOST = host;
        userIp = ip;
        webSocketSet.add(ip);
        System.out.println("当前连接数 : " + webSocketSet.size());

        this.session = webShell;
        JSch jsch = new JSch();
        com.jcraft.jsch.Session session = jsch.getSession(USER, HOST, DEFAULT_SSH_PORT);
        session.setPassword(PASSWORD);
        session.setConfig(STRICT_HOST_KEY_CHECKING, NO);
        UserInfo userInfo = new ShellInfo();
        session.setUserInfo(userInfo);
        session.connect(SESSION_TIMEOUT);
        channel = session.openChannel(JSCH_TYPE);
        try {
            in = channel.getInputStream();
            out = channel.getOutputStream();
        } catch (Exception e) {
            logger.error("实例化远程输出输出流发生错误！{}", new Object[]{e.getMessage()});
        }
        reader = new IoReader(this, THREAD_NAME);
        channel.connect(CHANNEL_TIMEOUT);
    }

    /**
     * 连接关闭调用的方法
     */
    @OnClose
    public void onClose() {
        try {
            destory();
        } catch (Exception e) {
            logger.error("释放资源发生错误！{}", new Object[]{e.getMessage()});
        }
    }

    /**
     * 收到客户端消息后调用的方法
     *
     * @param message 客户端发送过来的消息
     */
    @OnMessage
    public void onMessage(String message) {
        try {
            out.write(message.getBytes(CHAESET_NAME));
            out.flush();
        } catch (Exception e) {
            logger.error("处理消息发生错误！{}", new Object[]{e.getMessage()});
        }
    }

    /**
     * 发生错误时调用
     */
    @OnError
    public void onError(Throwable error) {
        logger.error("发生未知错误！{}", new Object[]{error.getMessage()});
        try {
            destory();
        } catch (IOException e) {
            logger.error("释放资源发生错误！{}", new Object[]{e.getMessage()});
        }
    }

    public void write(byte[] message) throws IOException {
        if (!session.isOpen()) {
            return;
        }
        try {
            this.session.getBasicRemote().sendBinary(ByteBuffer.wrap(message));
        } catch (Exception e) {
            logger.error("处理消息发生错误！{}", new Object[]{e.getMessage()});
        }
    }

    public Session getSession() {
        return session;
    }

    public InputStream getInputStream() throws IOException {
        return in;
    }

    private void destory() throws IOException {
        webSocketSet.remove(userIp);
        System.out.println("当前连接数 : " + webSocketSet.size());
        if (reader != null) {
            reader.destory();
            reader = null;
        }
        if (channel != null) {
            channel.disconnect();
            channel = null;
        }
        if (session != null) {
            session.close();
            session = null;
        }
        if (in != null) {
            in.close();
        }
        if (out != null) {
            out.close();
        }
    }
}